Support for Azure workload identity in AKS and Arc clusters #141

agreaves-ms · 2025-12-11T21:03:36Z

Description

Adds Azure AKS Workload Identity support, allowing pods to authenticate to Azure Blob Storage without connection strings or managed secrets. Uses DefaultAzureCredential for token-based auth and updates Helm charts to support workload identity configuration.

Related to Issue #153

Changes

Helm charts

all charts (service, router, backend-operator, web-ui) now support ServiceAccount configuration:
- serviceAccount.create - toggle chart-managed ServiceAccount creation
- serviceAccount.name - specify pre-provisioned SA name
- serviceAccount.annotations - add workload identity annotations (e.g., azure.workload.identity/client-id)
added extraPodLabels to services for workload identity labels (azure.workload.identity/use: "true")
made imagePullSecrets optional when using workload identity

Azure storage backend

AzureBlobStorageClient supports both connection string and token-based auth
added DefaultAzureCredential support for workload identity
SAS token generation now uses user delegation keys instead of account keys
credential fields (access_key_id, access_key) are now optional

Credentials

BasicDataCredential fields are now optional
added get_access_key_value() for safe credential retrieval
updated callers to use new accessor pattern

PyInstaller

added azure.identity, msal, and msal_extensions to Azure hooks for CLI workload identity support

Tests

added test coverage for workload identity auth paths

Usage

serviceAccount:
  create: false
  name: "my-workload-identity-sa"
  annotations:
    azure.workload.identity/client-id: "<client-id>"

services:
  service:
    extraPodLabels:
      azure.workload.identity/use: "true"

Checklist

I am familiar with the Contributing Guidelines.
New or existing tests cover these changes.
The documentation is up to date with these changes.

fernandol-nvidia · 2025-12-12T18:24:44Z

This is awesome, @agreaves-ms! Thank you so much for looking into this.

The team will be going through this change and thinking through what else we might need to change. Some high level thoughts is we probably need to give you more "types" of Credentials (e.g. WorkloadIdentityCredential, DataAccessCredential, ConnectionStringCredential, etc).

Another thing we will look at as a team is changes necessary to our service and osmo-ctrl to support this. Please stay tuned! And again, really appreciate the great work so far!

agreaves-ms · 2025-12-12T18:35:04Z

This is awesome, @agreaves-ms! Thank you so much for looking into this.

The team will be going through this change and thinking through what else we might need to change. Some high level thoughts is we probably need to give you more "types" of Credentials (e.g. WorkloadIdentityCredential, DataAccessCredential, ConnectionStringCredential, etc).

Another thing we will look at as a team is changes necessary to our service and osmo-ctrl to support this. Please stay tuned! And again, really appreciate the great work so far!

This is a great idea, I wanted to try and make a PR that was small and focused without refactoring to much of your codebase. Mainly to get the ideas through. Having a new credential type would be the most ideal, I'm needing to also alter the client_config.py to support DataSet upload and also DataSet download from both the osmo cli (if I'm a user running locally with az login done and I have the proper roles to be able to access the Storage Accoutn and Blob container) and then also from the osmo service kicking off the workload (needing to download the DataSet input data or upload the output data to the DataSet).

I'll update this PR with that change as well as a separate commit and then point you at it.

I was trying to think through a more agnostic naming scheme for this since Workload Identity is really an Azure and GCP term whereas AWS uses something like IRSA. I was thinking of going with EnvironmentAuthCredential or something similar... Anyways, you'll see that naming in this commit that I have.

Also, feel free to not take this PR, and instead implement the logic however way makes the most sense for your teams architecture that you've planned out. Thank you again for looking into this!

agreaves-ms · 2025-12-13T02:01:41Z

@fernandol-nvidia I've pushed an extra commit to further address workload identity support for input and output datasets, both on upload from the osmo-cli and then running in the workflow. Please take a look at that as well when considering implementation. I've now fully tested this in my own environment and I'm able to use my authentication with az login before doing an osmo workflow submit and it uses my credentials and access to the storage account to upload the dataset, then in the running pod for the workflow it's using the workload identity that's there to pull in the dataset in the blob.

Feel free to use this work however your team needs. I'm happy to further test any implementation that you all make available as well.

fernandol-nvidia · 2025-12-18T17:22:24Z

Hey @agreaves-ms, thanks for the update and apologies for the delay. I spent the past couple of days doing a refactor of how OSMO data operations work w.r.t to different credential strategies. I have a pending PR at #159.

The tl;dr is that we are differentiating DataCredential into StaticDataCredential (what we currently have) and DefaultDataCredential (which we delegate to the underlying SDK on credential resolution).

I believe this would give you the necessary abstractions/interface to add workload identity support for Azure. More specifically:

Defining the behavior at

OSMO/src/lib/data/storage/backends/azure.py

Lines 274 to 287 in 682dbf3

    
           def create_client(data_cred: credentials.DataCredential) -> blob.BlobServiceClient: 
        
               """ 
        
               Creates a new Azure Blob Storage client. 
        
               """ 
        
               match data_cred: 
        
                   case credentials.StaticDataCredential(): 
        
                       return blob.BlobServiceClient.from_connection_string( 
        
                           conn_str=data_cred.access_key.get_secret_value(), 
        
                       ) 
        
                   case credentials.DefaultDataCredential(): 
        
                       raise NotImplementedError( 
        
                           'Default data credentials are not supported yet') 
        
                   case _ as unreachable: 
        
                       assert_never(unreachable)

And overriding this method on how to decide if workload identity should be available for a particular backend:

OSMO/src/lib/data/storage/backends/common.py

Lines 279 to 297 in 682dbf3

    
               @functools.cached_property 
        
               def resolved_data_credential(self) -> credentials.DataCredential: 
        
                   """ 
        
                   Resolve the data credential for the storage backend. 
        
                   Subclasses should override this method to provide custom credential resolution logic. 
        
                   Returns: 
        
                       The resolved data credential. 
        
                   Raises: 
        
                       OSMOCredentialError: If the data credential is not found. 
        
                   """ 
        
                   data_cred = credentials.get_static_data_credential_from_config(self.profile) 
        
                   if data_cred is None: 
        
                       raise osmo_errors.OSMOCredentialError( 
        
                           f'Data credential not found for {self.profile}') 
        
                   return data_cred

Hopefully, this frees you from having to update so many unrelated files just to get Azure workload identity working :)

Please feel free to add any feedback/suggestions that would be helpful for you. I will keep you updated as things progress and can help you rebase this PR to absorb the new changes.

agreaves-ms · 2026-01-12T06:11:03Z

@fernandol-nvidia I've update this PR with the changes that have since been added. I needed to make some additional changes to parts that would allow for DataCredential instead of StaticDataCredential (only) as well.

Please review, I built this and verified all the azure workload identity, RBAC, and local auth works with these changes. Feel free to take over this PR as well as needed.

fernandol-nvidia

Happy 2026 :)

Nice, this looks great! I've added some minor comments which will help slim down the changes necessary.

Given that you've already tested this, I feel pretty confident with the PR. Once comments are addressed, I will approve!

src/utils/connectors/postgres.py

src/utils/job/task.py

src/lib/data/storage/backends/azure.py

- implement service account creation with customizable name and annotations - enhance service templates to support extra pod labels for various services - update Azure backend to utilize DefaultAzureCredential for authentication - add tests for Azure credential extraction and client creation

…Storage - add function to extract AccountKey from connection string - update AzureBlobStorageClient to handle different credential types

…n strings

…TaskGroup - change StaticDataCredential to DataCredential in get_all_data_creds method - update fetch_creds function signature to use DataCredential

…d account URL - remove deprecated storage account extraction function - modify create_client to accept storage_account and account_url parameters - update AzureBlobStorageClientFactory to use new parameters - adjust tests to reflect changes in client creation 🔒 - Generated by Copilot

…ent function 🔧 - Generated by Copilot

…eation 🔧 - Generated by Copilot

agreaves-ms · 2026-01-24T08:14:45Z

@fernandol-nvidia I finally had some time to be able to make these changes and redeploy to test them. Please review and let me know if any additional changes are needed.

fernandol-nvidia

Looks great to me! Thank you so much for making the changes.

* allow flexible squid proxy replicas (#241) * allow flexible squid proxy replicas * fix * Efficient Workflow Cleanup through Using Async Operations for Log Migration (#167) * Improving Performance for Uploading Workflow Artifacts in Worker Jobs * Cleanup * Add progress writing after upload * Add dependency in Bazel BUILD * Add type to mypy requirements * Update mypy requirements * Add to mypy_cli BUILD * Fix lint * Comment * Use constant to define semaphor and storage client executor count * #244 - Use last login url if url is not specified (#245) * Use last login url if url is not specified * print message * Cannot select any text inside modals or slideouts (#248) * Video html element not changin when selecting different video files in the UI for OSMO dataset (#249) * sync-feature-branches: fix no conflict case, allow single branch to be synced (#252) * Fix sync-feature-branches with no merge conflicts * Allow a single branch to be specified for sync-feature-branches * Perform operations as OSMO CI Bot * Add external label when the PR is created * extract issue number * add test cases (#247) * Allow PR checks to run on release branches (#264) * Database Pooling in Postgres Singleton Across Services (#251) * Initial commit for database pooling * Update set_session * Fix lint * Update PostgresConnector to have semaphor to control connections * Lint fix * Fix number of maxconn for test * Address comments * Add Go Postgres utils (#272) * #148 - Auth Project Design Documents (#165) * add args to postgres (#282) * #267 - cloud deployment scripts (#268) * script to create azure resources and deploy * Remove auto-generated values files from tracking - Added .gitignore to ignore values/, *.env files - Removed values/*.yaml files from git (auto-generated during deployment) * add aws script * add aws script * add copyright * update copyright * Support for Azure workload identity in AKS and Arc clusters (#141) * feat(src): add Azure service account and extra pod labels configuration - implement service account creation with customizable name and annotations - enhance service templates to support extra pod labels for various services - update Azure backend to utilize DefaultAzureCredential for authentication - add tests for Azure credential extraction and client creation * feat(src): extract account key from connection string for Azure Blob Storage - add function to extract AccountKey from connection string - update AzureBlobStorageClient to handle different credential types * feat(test): add tests for account key extraction from Azure connection strings * chore: clean up linting issues for tests * refactor(src): update data credential types in PostgresConnector and TaskGroup - change StaticDataCredential to DataCredential in get_all_data_creds method - update fetch_creds function signature to use DataCredential * feat(src): update Azure client creation to include storage account and account URL - remove deprecated storage account extraction function - modify create_client to accept storage_account and account_url parameters - update AzureBlobStorageClientFactory to use new parameters - adjust tests to reflect changes in client creation 🔒 - Generated by Copilot * refactor(src): mark storage_account parameter as unused in create_client function 🔧 - Generated by Copilot * refactor(src): remove unused storage_account parameter from client creation 🔧 - Generated by Copilot * Fix conflicts --------- Co-authored-by: Vivian Pan <[email protected]> Co-authored-by: ethany-nv <[email protected]> Co-authored-by: RyaliNvidia <[email protected]> Co-authored-by: patclarknvidia <[email protected]> Co-authored-by: Ethan Look-Potts <[email protected]> Co-authored-by: xutongNV <[email protected]> Co-authored-by: Allen Greaves <[email protected]>

* allow flexible squid proxy replicas (#241) * allow flexible squid proxy replicas * fix * Efficient Workflow Cleanup through Using Async Operations for Log Migration (#167) * Improving Performance for Uploading Workflow Artifacts in Worker Jobs * Cleanup * Add progress writing after upload * Add dependency in Bazel BUILD * Add type to mypy requirements * Update mypy requirements * Add to mypy_cli BUILD * Fix lint * Comment * Use constant to define semaphor and storage client executor count * #244 - Use last login url if url is not specified (#245) * Use last login url if url is not specified * print message * Cannot select any text inside modals or slideouts (#248) * Video html element not changin when selecting different video files in the UI for OSMO dataset (#249) * sync-feature-branches: fix no conflict case, allow single branch to be synced (#252) * Fix sync-feature-branches with no merge conflicts * Allow a single branch to be specified for sync-feature-branches * Perform operations as OSMO CI Bot * Add external label when the PR is created * extract issue number * add test cases (#247) * Allow PR checks to run on release branches (#264) * Database Pooling in Postgres Singleton Across Services (#251) * Initial commit for database pooling * Update set_session * Fix lint * Update PostgresConnector to have semaphor to control connections * Lint fix * Fix number of maxconn for test * Address comments * Add Go Postgres utils (#272) * #148 - Auth Project Design Documents (#165) * add args to postgres (#282) * #267 - cloud deployment scripts (#268) * script to create azure resources and deploy * Remove auto-generated values files from tracking - Added .gitignore to ignore values/, *.env files - Removed values/*.yaml files from git (auto-generated during deployment) * add aws script * add aws script * add copyright * update copyright * Support for Azure workload identity in AKS and Arc clusters (#141) * feat(src): add Azure service account and extra pod labels configuration - implement service account creation with customizable name and annotations - enhance service templates to support extra pod labels for various services - update Azure backend to utilize DefaultAzureCredential for authentication - add tests for Azure credential extraction and client creation * feat(src): extract account key from connection string for Azure Blob Storage - add function to extract AccountKey from connection string - update AzureBlobStorageClient to handle different credential types * feat(test): add tests for account key extraction from Azure connection strings * chore: clean up linting issues for tests * refactor(src): update data credential types in PostgresConnector and TaskGroup - change StaticDataCredential to DataCredential in get_all_data_creds method - update fetch_creds function signature to use DataCredential * feat(src): update Azure client creation to include storage account and account URL - remove deprecated storage account extraction function - modify create_client to accept storage_account and account_url parameters - update AzureBlobStorageClientFactory to use new parameters - adjust tests to reflect changes in client creation 🔒 - Generated by Copilot * refactor(src): mark storage_account parameter as unused in create_client function 🔧 - Generated by Copilot * refactor(src): remove unused storage_account parameter from client creation 🔧 - Generated by Copilot * Add new project proposal to describe nvlink + topology aware scheduling (#211) * Add new project proposal to describe nvlink + topology aware scheduling * Split design into two docs * Finish docs and add some updates from feedback * Add some open items * OSMO-6044: Application error when closing Task Details after switching Events view from Task to Workflow (#315) * add redis utlis, update postgres utils (#313) * add redis utlis, update postgres utils * add deps * Fix missing seperator in the test runner roles (#320) * fix * remove * fix --------- Co-authored-by: Vivian Pan <[email protected]> Co-authored-by: ethany-nv <[email protected]> Co-authored-by: RyaliNvidia <[email protected]> Co-authored-by: patclarknvidia <[email protected]> Co-authored-by: Ethan Look-Potts <[email protected]> Co-authored-by: xutongNV <[email protected]> Co-authored-by: Allen Greaves <[email protected]> Co-authored-by: ecolternv <[email protected]> Co-authored-by: tdewanNvidia <[email protected]>

agreaves-ms requested a review from a team December 11, 2025 21:03

agreaves-ms force-pushed the feat/azure-aks-workload-identity branch from f2d996c to 6bd2877 Compare January 12, 2026 06:07

agreaves-ms mentioned this pull request Jan 12, 2026

feat: support identity updates from OSMO Azure-Samples/azure-nvidia-robotics-reference-architecture#40

Open

fernandol-nvidia reviewed Jan 13, 2026

View reviewed changes

src/utils/connectors/postgres.py Show resolved Hide resolved

src/utils/job/task.py Show resolved Hide resolved

src/lib/data/storage/backends/azure.py Outdated Show resolved Hide resolved

agreaves-ms added 6 commits January 23, 2026 07:27

feat(src): extract account key from connection string for Azure Blob …

20a992b

…Storage - add function to extract AccountKey from connection string - update AzureBlobStorageClient to handle different credential types

feat(test): add tests for account key extraction from Azure connectio…

3d2fa17

…n strings

chore: clean up linting issues for tests

b2091bb

refactor(src): update data credential types in PostgresConnector and …

de6bebe

…TaskGroup - change StaticDataCredential to DataCredential in get_all_data_creds method - update fetch_creds function signature to use DataCredential

agreaves-ms force-pushed the feat/azure-aks-workload-identity branch from 30edbb7 to d66d6a8 Compare January 23, 2026 11:00

agreaves-ms added 2 commits January 23, 2026 14:13

refactor(src): mark storage_account parameter as unused in create_cli…

e148b0b

…ent function 🔧 - Generated by Copilot

refactor(src): remove unused storage_account parameter from client cr…

fb4e144

…eation 🔧 - Generated by Copilot

fernandol-nvidia approved these changes Jan 26, 2026

View reviewed changes

vvnpn-nv approved these changes Jan 26, 2026

View reviewed changes

fernandol-nvidia merged commit f8ca85b into NVIDIA:main Jan 26, 2026
9 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support for Azure workload identity in AKS and Arc clusters #141

Support for Azure workload identity in AKS and Arc clusters #141

Uh oh!

agreaves-ms commented Dec 11, 2025 •

edited

Loading

Uh oh!

fernandol-nvidia commented Dec 12, 2025

Uh oh!

agreaves-ms commented Dec 12, 2025

Uh oh!

agreaves-ms commented Dec 13, 2025

Uh oh!

fernandol-nvidia commented Dec 18, 2025

Uh oh!

agreaves-ms commented Jan 12, 2026 •

edited

Loading

Uh oh!

fernandol-nvidia left a comment •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

agreaves-ms commented Jan 24, 2026

Uh oh!

fernandol-nvidia left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Support for Azure workload identity in AKS and Arc clusters #141

Support for Azure workload identity in AKS and Arc clusters #141

Uh oh!

Conversation

agreaves-ms commented Dec 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Changes

Helm charts

Azure storage backend

Credentials

PyInstaller

Tests

Usage

Checklist

Uh oh!

fernandol-nvidia commented Dec 12, 2025

Uh oh!

agreaves-ms commented Dec 12, 2025

Uh oh!

agreaves-ms commented Dec 13, 2025

Uh oh!

fernandol-nvidia commented Dec 18, 2025

Uh oh!

agreaves-ms commented Jan 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

fernandol-nvidia left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

agreaves-ms commented Jan 24, 2026

Uh oh!

fernandol-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

agreaves-ms commented Dec 11, 2025 •

edited

Loading

agreaves-ms commented Jan 12, 2026 •

edited

Loading

fernandol-nvidia left a comment •

edited

Loading